#include <functional>
#include <rclcpp/rclcpp.hpp>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2/LinearMath/Quaternion.h>

#include <nav_msgs/msg/odometry.hpp>
#include <sensor_msgs/msg/imu.hpp>
#include <std_msgs/msg/float64.hpp>

class Odom : public rclcpp::Node
{
  private:
    rclcpp::Publisher<nav_msgs::msg::Odometry>::SharedPtr odom_pub;
    rclcpp::Subscription<sensor_msgs::msg::Imu>::SharedPtr imu_sub;
    rclcpp::Subscription<std_msgs::msg::Float64>::SharedPtr left_sub;
    rclcpp::Subscription<std_msgs::msg::Float64>::SharedPtr right_sub;
    rclcpp::TimerBase::SharedPtr odom_timer;

    double leftVec = 0, rightVec = 0;
    tf2::Quaternion rotation = tf2::Quaternion();
    double angularVelocity = 0;
    double x = 0, y = 0;
    rclcpp::Time lastTime = rclcpp::Time(0);

  public:
    Odom() : Node("odom")
    {
        RCLCPP_INFO(this->get_logger(), "odom node start");
        lastTime = rclcpp::Time(0);
        imu_sub = create_subscription<sensor_msgs::msg::Imu>(
            "imu", 1,
            [this](sensor_msgs::msg::Imu::SharedPtr msg)
            {
                rotation =
                    tf2::Quaternion(msg->orientation.x, msg->orientation.y,
                                    msg->orientation.z, msg->orientation.w);
                rotation.normalize();
                angularVelocity = msg->angular_velocity.z;

                rclcpp::Time now = rclcpp::Time(msg->header.stamp);
                if (lastTime.nanoseconds() == 0)
                    lastTime = now;
                double dt = (now - lastTime).seconds();
                lastTime = now;

                tf2::Matrix3x3 m(rotation);
                double roll, pitch, yaw;
                m.getRPY(roll, pitch, yaw);
                double vec = (rightVec + leftVec) / 2;
                x += vec * cos(yaw) * dt;
                y += vec * sin(yaw) * dt;
                // RCLCPP_INFO(get_logger(), "x: %f, y: %f, vec:%f, dt: %f", x,
                // y, vec, dt);
            });
        left_sub = create_subscription<std_msgs::msg::Float64>(
            "left_speed", 1, [this](std_msgs::msg::Float64::SharedPtr msg)
            { leftVec = msg->data; });
        right_sub = create_subscription<std_msgs::msg::Float64>(
            "right_speed", 1, [this](std_msgs::msg::Float64::SharedPtr msg)
            { rightVec = msg->data; });
        odom_pub = this->create_publisher<nav_msgs::msg::Odometry>("odom", 10);
        odom_timer = this->create_wall_timer(
            std::chrono::milliseconds(10), std::bind(&Odom::PublishOdom, this));
    }

    void PublishOdom()
    {
        nav_msgs::msg::Odometry msg;
        msg.header.stamp = this->now();
        msg.header.frame_id = "map";
        msg.child_frame_id = "footprint";
        msg.pose.pose.position.x = x;
        msg.pose.pose.position.y = y;
        msg.pose.pose.position.z = 0;
        rotation.normalize();
        msg.pose.pose.orientation.x = rotation.x();
        msg.pose.pose.orientation.y = rotation.y();
        msg.pose.pose.orientation.z = rotation.z();
        msg.pose.pose.orientation.w = rotation.w();
        msg.twist.twist.linear.x = (leftVec + rightVec) / 2;
        msg.twist.twist.linear.y = 0;
        msg.twist.twist.linear.z = 0;
        msg.twist.twist.angular.x = 0;
        msg.twist.twist.angular.y = 0;
        msg.twist.twist.angular.z = angularVelocity;
        odom_pub->publish(msg);
    }
};

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<Odom>());
    rclcpp::shutdown();
    return 0;
}